from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from Crypto.Util.Padding import pad, unpad
import base64

from app.core.cons import Cons

# 密钥必须是16、24或32字节长
SECRET_KEY = b'abcesdkfjhsfsdlj'

class AESCrypto:
    @staticmethod
    def encrypt(data) -> str:
        # 生成随机IV
        iv = get_random_bytes(16)
        # 创建AES加密器
        cipher = AES.new(SECRET_KEY, AES.MODE_CBC, iv)
        # 对数据进行填充并加密
        # 检查数据类型，如果是字符串则先编码，否则直接使用
        if isinstance(data, str):
            padded_data = pad(data.encode(), AES.block_size)
        elif isinstance(data, bytes):
            padded_data = pad(data, AES.block_size)
        else:
            raise TypeError("Data must be str or bytes")
        encrypted_data = cipher.encrypt(padded_data)
        # 将IV和加密数据合并并进行base64编码
        combined = iv + encrypted_data
        return base64.b64encode(combined).decode(Cons.UTF_8)
    
    @staticmethod
    def decrypt(encrypted_data: str) -> str:
        # 解码base64数据
        combined = base64.b64decode(encrypted_data)
        # 提取IV和加密数据
        iv = combined[:16]
        encrypted = combined[16:]
        # 创建AES解密器
        cipher = AES.new(SECRET_KEY, AES.MODE_CBC, iv)
        # 解密并去除填充
        decrypted = cipher.decrypt(encrypted)
        unpadded = unpad(decrypted, AES.block_size)
        return unpadded.decode(Cons.UTF_8)
